#!/usr/bin/python
# Copyright (C) 2003, Charles Wang.
# Author:  Charles Wang <charles@linux.net.cn>
# All rights reserved.
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the Free
# Software Foundation; either version 2, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANT; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE. See the GNU General Public LIcense for more
# details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, 59 Temple
# Place - Suite 330, Boston, MA 02111-1307, USA.

import os
import sys
import kudzu
import parted

all_harddisks = {}

def device_probe_all():
    global all_harddisks
    result = []
    # The following commented code is implemented by kudzu.
    devlist = map(lambda hd: parted.PedDevice.get(os.path.join('/dev', hd.device)),
                  kudzu.probe(kudzu.CLASS_HD,
                              kudzu.BUS_IDE | kudzu.BUS_SCSI | kudzu.BUS_MISC,
                              kudzu.PROBE_ALL))
    if devlist:
        for dev in devlist:
            newdisklabel = None
            try:
                disk = parted.PedDisk.new(dev)
            except parted.error:
                # For the disk without any disk label, create it.
                print 'Warning: parted.error in device_probe_all()'
                newdisklabel = 'y'
                disk = dev.disk_new_fresh(parted.disk_type_get('msdos'))
            model = dev.model
            result.append((dev.path, dev.length, model))
            all_harddisks[dev.path] = (dev, disk, newdisklabel)
    return  result

def get_all_partitions(devpath):
    def part2result(part):
        flags = []
        avaflags = []
        label = 'N/A'
        if part.is_active() != 0:
            for f in range(parted.PARTITION_FIRST_FLAG,
                           parted.PARTITION_LAST_FLAG + 1):
                if part.is_flag_available(f):
                    avaflags.append(f)
                    if part.get_flag(f):
                        flags.append(f)
            if part.disk.type.check_feature(parted.DISK_TYPE_PARTITION_NAME):
                label = part.get_name()
        if part.fs_type:
            fs_type_name = part.fs_type.name
        else:
            fs_type_name = 'N/A'
        return (part.num, flags, part.type, part.geom.length,
                fs_type_name, label, part.geom.start, part.geom.end, avaflags)

    result = []
    if all_harddisks.has_key(devpath):
        disk = all_harddisks[devpath][1]
        if disk:
            part = disk.next_partition()
            while part:
                if part.type & parted.PARTITION_METADATA == 0:
                    result.append(part2result(part))
                part = disk.next_partition(part)
    return  result


print '----device_probe_all() result----'
print device_probe_all()
print '----get_all_partitions() result----'
for devpath in all_harddisks.keys():
    print 'Device is %s' % devpath
    for item in get_all_partitions(devpath):
        print item
